import os
import sys

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,4"


project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
print(project_dir_path)
sys.path.append(project_dir_path)

from omics_bert_label_encoding import MakeLabelsNhanes 
from omics_bert_label_encoding.label.check_label_files_nhanes import CheckLabelFilesNhanes
from omics_bert_label_encoding.label.processing_labels_nhanes import process_label_embeddings_nhanes


def main(dataset_path, save_path, dictionary_path, api_key, expected_lines, regenerate, model_name, tensor_parallel_size):
    print("--- Step 1: Build Labels ---")
    build_labels = MakeLabelsNhanes(dataset_path, save_path, dictionary_path, api_key)
    build_labels.main()

    print("\n--- Step 2: Check and Regenerate Labels ---")
    check_path = os.path.join(save_path, "labels")

    checker = CheckLabelFilesNhanes(
        check_path=check_path,
        expected_line_count=expected_lines,
        regenerate=regenerate,
        dataset_path=dataset_path,
        dictionary_path=dictionary_path,
        api_key=api_key
    )
    checker.main()

    print("\n--- Step 3: Process Label Embeddings ---")
    labels_dir = os.path.join(save_path, "labels")
    process_label_embeddings_nhanes(
        labels_json_path=labels_dir,
        output_dir=save_path,
        model_name=model_name,
        tensor_parallel_size=tensor_parallel_size
    )


if __name__ == '__main__':
    project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
    print(project_dir_path)
    sys.path.append(project_dir_path)

    dataset_path = "/home/xw/python3test/omics-bert-datasets/datasets/ready/nhanes/v1"
    save_path = os.path.join(project_dir_path, "embeddings_local/nhanes/v1-labels_embeddings")
    dictionary_path = os.path.join(project_dir_path, "datasets/raw/nhanes/v1/df_dictionary.xlsx")

    api_key = ""
    expected_lines = 12
    regenerate = True
    model_name = "Qwen/Qwen3-Embedding-8B"
    tensor_parallel_size = 2

    main(dataset_path, save_path, dictionary_path, api_key, expected_lines, regenerate, model_name, tensor_parallel_size)